import re
from typing import List, Dict

class QAEngine:
    """问答引擎，基于检索到的文档生成答案"""
    
    def __init__(self):
        self.max_context_length = 2000  # 最大上下文长度
    
    def generate_answer(self, question: str, relevant_docs: List[Dict]) -> str:
        """基于检索到的文档生成答案"""
        if not relevant_docs:
            return "抱歉，没有找到相关的文档内容来回答您的问题。"
        
        # 构建上下文
        context = self._build_context(relevant_docs)
        
        # 生成答案
        answer = self._generate_simple_answer(question, context)
        
        return answer
    
    def _build_context(self, docs: List[Dict]) -> str:
        """构建上下文"""
        context_parts = []
        total_length = 0
        
        for doc in docs:
            content = doc["content"]
            if total_length + len(content) <= self.max_context_length:
                context_parts.append(content)
                total_length += len(content)
            else:
                remaining_length = self.max_context_length - total_length
                if remaining_length > 50:  # 至少保留50个字符
                    context_parts.append(content[:remaining_length])
                break
        
        return "\n\n".join(context_parts)
    
    def _generate_simple_answer(self, question: str, context: str) -> str:
        """简单的答案生成（基于规则和关键词匹配）"""
        # 这里使用简单的规则生成答案，实际项目中可以集成LLM
        
        # 提取问题中的关键词
        keywords = self._extract_keywords(question)
        
        # 在上下文中查找相关信息
        relevant_sentences = self._find_relevant_sentences(context, keywords)
        
        if not relevant_sentences:
            return "根据现有文档内容，我无法找到确切的答案。建议您：\n1. 检查问题表述是否清晰\n2. 尝试使用不同的关键词提问\n3. 上传更多相关文档"
        
        # 生成答案
        answer = self._construct_answer(question, relevant_sentences)
        
        return answer
    
    def _extract_keywords(self, question: str) -> List[str]:
        """提取关键词"""
        # 移除标点符号和停用词
        import re
        
        # 移除标点
        question = re.sub(r'[^\w\s]', '', question)
        
        # 简单的关键词提取
        words = question.lower().split()
        
        # 过滤停用词和短词
        stop_words = {'什么', '怎么', '为什么', '哪里', '谁', '何时', '如何', 
                     '是', '的', '了', '在', '和', '有', '我', '你', '他', '她', '它',
                     '这', '那', '这些', '那些', '一个', '一些', '可以', '应该', '需要',
                     'what', 'how', 'why', 'where', 'who', 'when', 'which',
                     'is', 'are', 'was', 'were', 'the', 'a', 'an', 'and', 'or', 'but',
                     'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by', 'from'}
        
        keywords = [word for word in words if len(word) > 1 and word not in stop_words]
        
        return keywords[:5]  # 最多返回5个关键词
    
    def _find_relevant_sentences(self, context: str, keywords: List[str]) -> List[str]:
        """查找相关句子"""
        # 分割句子
        sentences = re.split(r'[。！？\n.!?]', context)
        
        relevant_sentences = []
        
        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue
            
            # 计算关键词匹配度
            score = 0
            for keyword in keywords:
                if keyword in sentence.lower():
                    score += 1
            
            # 如果匹配度足够高，添加到结果中
            if score > 0 and len(sentence) > 10:  # 至少10个字符
                relevant_sentences.append((sentence, score))
        
        # 按匹配度排序
        relevant_sentences.sort(key=lambda x: x[1], reverse=True)
        
        # 返回前3个最相关的句子
        return [sentence for sentence, score in relevant_sentences[:3]]
    
    def _construct_answer(self, question: str, sentences: List[str]) -> str:
        """构建答案"""
        if not sentences:
            return "抱歉，没有找到确切的答案。"
        
        # 简单的答案构建
        answer = "根据文档内容，我找到以下相关信息：\n\n"
        
        for i, sentence in enumerate(sentences, 1):
            answer += f"{i}. {sentence}"
            if not sentence.endswith(('。', '！', '？', '.', '!', '?')):
                answer += '。'
            answer += '\n\n'
        
        # 添加总结
        answer += "以上信息应该能回答您的问题。如果您需要更详细的解释，请告诉我。"
        
        return answer
